import cv2
import numpy as np
import pickle
import matplotlib.pyplot as plt

# 读取模型
def load_model(model_path):
    with open(model_path, 'rb') as f:
        clf = pickle.load(f)
    return clf

# 创建特征并预测
def create_features_and_labels(image):
    # 定义滤波器
    sobel_x = cv2.Sobel(image, cv2.CV_64F, 1, 0, ksize=3)
    sobel_y = cv2.Sobel(image, cv2.CV_64F, 0, 1, ksize=3)
    sobel_edges = np.sqrt(sobel_x**2 + sobel_y**2)

    # 使用Canny边缘检测
    canny_edges = cv2.Canny(image, 50, 150)

    # 将图像和边缘检测结果展平成一维数组
    X_image = image.reshape(-1, 1)
    X_sobel_edges = sobel_edges.reshape(-1, 1)
    X_canny_edges = canny_edges.reshape(-1, 1)

    # 合并新的特征列
    X_combined = np.concatenate((X_image, X_sobel_edges, X_canny_edges), axis=1)

    return X_combined

def predict_segmentation(clf, image):
    X = create_features_and_labels(image)  # 使用与训练时相同的create_features函数
    y_pred = clf.predict(X)
    return y_pred.reshape(image.shape)  # 将预测结果重塑为原始图像的形状

# 主函数
def main():
    # 加载模型
    clf = load_model('clf.pkl')
    
    # 读取新图像
    image2 = cv2.imread('C:/Users/86178/Desktop/Segment_sandstone/sandstone_2.tif', cv2.IMREAD_GRAYSCALE)
    label2 = cv2.imread('C:/Users/86178/Desktop/Segment_sandstone/sandstone_2_segment.tif', cv2.IMREAD_GRAYSCALE)
    
    # 预测分割
    predicted_segment = predict_segmentation(clf, image2)
    
    # 计算准确率（这里简单地将预测结果与真实标签进行比较）
    accuracy = np.sum(predicted_segment == label2) / np.prod(label2.shape)
    print(f"准确率: {accuracy:.4f}")
    
    # 显示图像
    plt.figure(figsize=(10, 5))
    
    plt.subplot(1, 3, 1)
    plt.title('Original Image')
    plt.imshow(image2, cmap='gray')
    
    plt.subplot(1, 3, 2)
    plt.title('True Segmentation')
    plt.imshow(label2, cmap='gray')
    
    plt.subplot(1, 3, 3)
    plt.title('Predicted Segmentation')
    plt.imshow(predicted_segment, cmap='gray')
    
    plt.show()

if __name__ == "__main__":
    main()